#include "SpectraSTDenoiser.hpp"
#include <iostream>
#include <algorithm>
#include <math.h>

using namespace std;

SpectraSTDenoiser::SpectraSTDenoiser() :
  m_signalWithCNI(100, 0),
  m_noiseWithCNI(100, 0),
  m_signalWithMzPos(100, 0),
  m_noiseWithMzPos(100, 0),
  m_ready(false) {

  m_sisters.push_back(-2.016);
  m_sisters.push_back(-1.008);
  m_sisters.push_back(1.008);
  m_sisters.push_back(2.016);
  m_sisters.push_back(9.005285);
  m_sisters.push_back(17.0027);
  m_sisters.push_back(18.01057);
  m_sisters.push_back(27.9949);
  m_sisters.push_back(57.02146);
  m_sisters.push_back(71.03711);
  m_sisters.push_back(87.03202);
  m_sisters.push_back(97.05276);
  m_sisters.push_back(99.06841);
  m_sisters.push_back(101.0476);
  m_sisters.push_back(103.0091);
  m_sisters.push_back(113.084);
  m_sisters.push_back(114.0429);
  m_sisters.push_back(115.0269);
  m_sisters.push_back(128.0585);
  m_sisters.push_back(128.0949);
  m_sisters.push_back(129.0425);
  m_sisters.push_back(131.0404);
  m_sisters.push_back(137.0589);
  m_sisters.push_back(147.0684);
  m_sisters.push_back(156.1011);
  m_sisters.push_back(163.0633);
  m_sisters.push_back(186.0793);
    
  m_signalWithSister.assign(m_sisters.size(), 0);
  m_noiseWithSister.assign(m_sisters.size(), 0);
  
}

SpectraSTDenoiser::~SpectraSTDenoiser() {

  
}

void SpectraSTDenoiser::useDefault() {
  

  
  m_ready = true;
  
}



void SpectraSTDenoiser::addTrainingSpectrum(SpectraSTPeakList* pl) {

  if (m_ready) return;
  
  float cumInten = 0.0;

  for (unsigned int rank = pl->getNumPeaks(); rank >= 1; rank--) {
    
    Peak p;
    pl->getNthLargestPeak(rank, p);
    
    // ignore all precursor peaks?
    if (pl->isNearPrecursor(p.mz)) {
      continue;
    }

    // calculate CNI
    cumInten += p.intensity;
    float cni = cumInten / pl->getTotalIonCurrent();
    int cniBin = (int)(cni * 100.0);
    
    // calculate m/z position
    int mzPosBin = (int)(p.mz / pl->getMaxMz() * 100.0);

    // find complement
    double precursorMass = pl->getParentMz() * (double)(pl->getParentCharge());
    float complement = pl->findPeak(precursorMass - p.mz, 0.25);
    bool hasComplement = false;
    if (complement >= 0.0) hasComplement = true;
    
    if (p.info == "S") {
      m_numSignal++;
      m_signalWithCNI[cniBin]++;
      m_signalWithMzPos[mzPosBin]++; 
      if (hasComplement) {
        m_signalWithComplement++;
      } 
	
    } else {
      m_numNoise++;
      m_noiseWithCNI[cniBin]++;
      m_noiseWithMzPos[mzPosBin]++;
      if (hasComplement) {
        m_noiseWithComplement++;
      } 
    }
    
    // find sisters
    for (unsigned int sisIndex = 0; sisIndex < (unsigned int)(m_sisters.size()); sisIndex++) {
      double mzDiff = m_sisters[sisIndex];
      if (pl->findPeak(p.mz - mzDiff, 0.25) >= 0.0) {
        if (p.info == "S") {
	  m_signalWithSister[sisIndex]++;
	} else {
	  m_noiseWithSister[sisIndex]++;
	}
      }
    }

    
    
  } // END for all peaks 
  
  
}

void SpectraSTDenoiser::generateBayesianModel() {
 
  if (m_ready) return;
  
  
  // calculate the conditional probabilities or odds ratio based on the signal/noise counts
  
  // for example: this is the lg( P(S|CNI = cniBin) / P(N|CNI = cniBin) )
  // m_logOddsCNI[cniBin] = log(((double)m_signalCNI[cniBin] / (double)m_numSignals) / ((double)m_noiseCNI[cniBin] / (double)m_numNoise));
  // ...

/*
  for (unsigned int index = 0; index <= 100; index++) {
    m_logOddsWithCNI[index] = log(m_signalWithCNI[index]) + log(m_numNoise) - log(m_noiseWithCNI[index])- log(m_numSignal);
    m_logOddsWithMzPos[index] = log(m_signalWithMzPos[index]) + log(m_numNoise) - log(m_noiseWithMzPos[index]) - log(m_numSignal); 
  }

  for (unsigned int sisIndex = 0; sisIndex < (unsigned int)(m_sisters.size()); sisIndex++) {
    m_logOddsWithSister[sisIndex] = log(m_signalWithSister[sisIndex]) + log(m_numNoise) - log(m_noiseWithSister[sisIndex]) - log(m_numSignal);

    m_logOddsWithoutSister[sisIndex] = log(m_numSignal - m_signalWithSister[sisIndex]) + log(m_numNoise) - log(m_numNoise - m_noiseWithSister[sisIndex]) - log(m_numSignal);
}
*/
  cerr << m_numNoise << ',' << m_numSignal << endl;
  cerr << m_noiseWithComplement << ',' << m_signalWithComplement << endl;
  
  m_logOddsWithComplement = log(m_signalWithComplement) + log(m_numNoise) - log(m_noiseWithComplement) - log(m_numSignal);

  m_logOddsWithoutComplement = log(m_numSignal - m_signalWithComplement) + log(m_numNoise) - log(m_numNoise - m_noiseWithComplement) - log(m_numSignal);

  m_ready = true;

}

void SpectraSTDenoiser::printModel() {
  
  cout << "MODEL:" << endl;
  cout << "logOdds(with complement) = " << m_logOddsWithComplement << endl;
  cout << "logOdds(without complement) = " << m_logOddsWithoutComplement << endl;
  
  
}

void SpectraSTDenoiser::filter(SpectraSTPeakList* pl, unsigned int maxNumPeaks, double minSignalProb) {
  
  if (!m_ready) return;
  // calculate signal prob for all peaks 
  
  vector<pair<double, Peak> > peakOdds; // vector of (logOdds, Peak) pairs
  
  float cumInten = 0.0;
  for (unsigned int rank = pl->getNumPeaks(); rank >= 1; rank--) {
    
    Peak p;
    pl->getNthLargestPeak(rank, p);
    
    // ignore all precursor peaks?
    if (pl->isNearPrecursor(p.mz)) {
      continue;
    }

    // calculate CNI
    cumInten += p.intensity;
    float cni = cumInten / pl->getTotalIonCurrent();
    int cniBin = (int)(cni * 100.0);
    
    // calculate m/z position
    int mzPosBin = (int)(p.mz / pl->getMaxMz() * 100.0);

    // find complement
    double precursorMass = pl->getParentMz() * (double)(pl->getParentCharge());
    float complement = pl->findPeak(precursorMass - p.mz, 0.25);
    bool hasComplement = false;
    if (complement >= 0.0) hasComplement = true;
    
    double logOdds = 0.0;
    
    logOdds += m_logOddsWithCNI[cniBin];
    logOdds += m_logOddsWithMzPos[mzPosBin];
    if (hasComplement) {
      logOdds += m_logOddsWithComplement;
    } else {
      logOdds += m_logOddsWithoutComplement;
    }
      
    // find sisters
    for (unsigned int sisIndex = 0; sisIndex < (unsigned int)(m_sisters.size()); sisIndex++) {
      double mzDiff = m_sisters[sisIndex];
      if (pl->findPeak(p.mz - mzDiff, 0.25) >= 0.0) {
 	logOdds += m_logOddsWithSister[sisIndex];
      } else {
        logOdds += m_logOddsWithoutSister[sisIndex];
      }
    }

    pair<double, Peak> op;
    op.first = logOdds;
    op.second = p;
    
    peakOdds.push_back(op);
  } // END for all peaks 
  
  
  // sort all peaks by odds
  sort(peakOdds.begin(), peakOdds.end(), SpectraSTDenoiser::sortPeaksByOddsDesc);
  
  // insert peaks back into new peak list
  pl->m_peaks.clear();
  
  double prior = 0.5;  
  if (minSignalProb > 0.0) {
    prior = predictPrior(pl);
  }
  
  double logOddsPrior = log(prior / (1.0 - prior));
  double logOddsCutoff = log(minSignalProb / (1.0 - minSignalProb));
  
  unsigned int numPeaksKept = 0;
  
  for (vector<pair<double, Peak> >::iterator i = peakOdds.begin(); i != peakOdds.end(); i++) {
    if (numPeaksKept >= maxNumPeaks) break;
    if (i->first + logOddsPrior < logOddsCutoff) break;
    
    pl->insert(i->second.mz, i->second.intensity, i->second.annotation, i->second.info);
    numPeaksKept++;
  }
  
  sort(pl->m_peaks.begin(), pl->m_peaks.end(), SpectraSTPeakList::sortPeaksByMzAsc);
  pl->m_isSortedByMz = true;
  
  if (pl->m_intensityRanked) {
    delete (pl->m_intensityRanked);
    pl->m_intensityRanked = NULL;
  }

  if (pl->m_peakMap) {
    delete (pl->m_peakMap);
    pl->m_peakMap = NULL;
  }

  
}

double SpectraSTDenoiser::predictPrior(SpectraSTPeakList* pl, bool useOverallPrior) {
  
  if (useOverallPrior) {
    return ((double)m_numSignal / (double)(m_numSignal + m_numNoise));
  }
  
  double prior = 0.2082;
  
  // ...
  unsigned int numPeaks = pl->getNumPeaks();
  double precursorMass = pl->getParentMz() * pl->getParentCharge();

  
  prior = 3.792e-01 + -3.788e-04 * numPeaks + -2.172e-05 * precursorMass;

  if(prior >= 1.0 || prior <= 0) prior = 0.2082;
  
  return (prior);
  
}

bool SpectraSTDenoiser::sortPeaksByOddsDesc(pair<double, Peak> a, pair<double, Peak> b) {
 
  return (a.first > b.first);
  
}
